package com.suppor.rpc.register.loadbalance;

import com.suppor.rpc.core.register.RegisterInfo;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * 基于JDK TreeMap 实现一致性Hash
 * 工程上需要优化为预加载，提升运行性能，这里为功能实现，暂不考虑线上情况
 */
public class ConsistentHashLoadBalance implements ServiceBalancer<RegisterInfo> {
    /**
     * 虚拟节点个数
     */
    private static final int VIRTUAL_NODE_NUM = 10;

    @Override
    public RegisterInfo select(List<RegisterInfo> services, int hashCode) {
        TreeMap<Integer, RegisterInfo> ring = new TreeMap<>();
        createConsistentHashRing(ring, services);
        return allocateNode(ring, hashCode);
    }

    private RegisterInfo allocateNode(TreeMap<Integer, RegisterInfo> ring, int hashCode) {
        Map.Entry<Integer, RegisterInfo> entry = ring.ceilingEntry(hashCode);
        if (entry == null) {
            entry = ring.firstEntry();
        }
        return entry.getValue();
    }

    private void createConsistentHashRing(TreeMap<Integer, RegisterInfo> ring, List<RegisterInfo> services) {
        for (RegisterInfo service : services) {
            for (int i = 0; i < VIRTUAL_NODE_NUM; i++) {
                int ringKey = (key(service) + "#" + i).hashCode();
                ring.put(ringKey, service);
            }
        }
    }

    public String key(RegisterInfo registerInfo) {
        return String.join(":", registerInfo.getServerAddr(), registerInfo.getPort().toString());
    }
}
